import sys
from pathlib import Path

base_path = str(Path(__file__).resolve().parent.parent)
sys.path.append(base_path)
from langchain_community.chat_message_histories import ChatMessageHistory
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables import RunnablePassthrough
from langchain_core.output_parsers import StrOutputParser
from create_llm import create_llm

# 初始化聊天历史
message_history = ChatMessageHistory()


def build_history(user_content):
    messages = message_history.messages
    # messages.append({"role": "user", "content": user_content})
    return messages


def save_history(input_text, result):
    message_history.add_user_message(input_text)
    message_history.add_ai_message(result)


prompt = ChatPromptTemplate.from_messages(
    [
        ("system", "你一名智能助手"),
        MessagesPlaceholder(variable_name="prompt_history"),
        ("user", "{user_content}"),
    ]
)


def log(val):
    print(val)
    print("------------------log---------------")
    return val


llm = create_llm()

chain = (
    RunnablePassthrough.assign(
        prompt_history=lambda x: build_history(x["user_content"])
    )
    | prompt
    | log
    | llm
    | StrOutputParser()
)


def get_input(prompt="\n请输入对话内容："):
    return input(prompt)


input_text = get_input()
while input_text != "exit":
    result = chain.invoke({"user_content": input_text})
    print(result)
    print("------------------result---------------")
    save_history(input_text, result)
    input_text = get_input()
